# -*- coding:utf-8 -*-
"""
@FileName  :etf_cta_strategy_shenshu.py
@Time      :2022/10/27 10:44
@Author    :fsksf

ETF 赎回示例
"""
import datetime
import time
from enum import Enum
from typing import Any, Callable, Dict
from vnpy.trader.utility import ArrayManager, BarGenerator
from vnpy.trader.object import BarData, TickData, OrderData, TradeData
# 注意这里引入的是ETFTemplate
from vnpy_ctastrategy.etf_template import ETFTemplate


class SSMode(Enum):
    DISCOUNT = "折价"
    PREMIUM = "溢价"
    NONE = "无操作"


class ETFBSStrategyShenShu(ETFTemplate):

    # 限定开仓交易时间，防止尾盘开仓
    trade_morning_start = datetime.time(hour=9, minute=32)      # 上午开始时间
    trade_morning_end = datetime.time(hour=11, minute=25)       # 上午结束时间
    trade_afternoon_start = datetime.time(hour=13, minute=32)   # 下午开始时间
    trade_afternoon_end = datetime.time(hour=14, minute=50)     # 下午结束时间

    # 启动策略时的配置项
    trade_basket = True
    discount_min_profit: float = 500                        # 折价最小利润
    premium_min_profit: float = 500                         # 溢价最小利润
    pre_ss_vol: int = 900000                                # 申赎最小单位
    per_order_vol: int = 1000000                            # 每次买卖ETF的最大数量，拆单，每次委托最高一百万(柜台规则)
    parameters = ['pre_ss_vol', "per_order_vol",
                  'discount_min_profit', 'premium_min_profit']

    def __init__(
        self,
        cta_engine: Any,
        strategy_name: str,
        vt_symbol: str,         # 篮子对应的ETF
        setting: dict,
    ):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.trade_basket = True
        self.tick_am = []
        self.am = ArrayManager(size=100)
        self.bg = BarGenerator(on_bar=self.on_bar)
        self.tick: TickData = None
        self.last_order_time = None
        self.mode: SSMode = SSMode.NONE                 # 当前模式：标识当前是折价交易、溢价交易状态

    def on_init(self):
        self.load_bar(1)

    def on_tick(self, tick: TickData):
        """
        这里之做bar的合成
        :param tick:
        :return:
        """
        self.tick = tick
        self.bg.update_tick(tick)
        if not self.trading:
            return
        self.handle_tick(tick)

    def handle_tick(self, tick):
        """
        tick级别的策略逻辑写在这里
        :param tick:
        :return:
        """
        # 1.是否有未完成的订单
        if not self.all_order_finished():
            if self.last_order_time and time.time() - self.last_order_time > 10:    # 委托撤单等待时间
                self.cancel_all()
            return
        # 2.下单后等待一段时间，保证trade已经返回
        if self.last_order_time and time.time() - self.last_order_time < 8:
            return

        # 3.获取瞬时利润
        # {'fund_code': '510500', 'fund_name': '南方中证500', 'excel_update_at': '20230920',
        # 'diff': '-247.34', 'update_at': '2024-01-29 10:54:05', 'discount': '-329.18',
        # 'premium': '-4477.64', 'dp_up_at': '2024-01-26 16:20:07'}
        moment_profit = self.cta_engine.get_moment_profit(self.vt_symbol[:6])
        if not moment_profit:
            self.print(f"策略 {self.strategy_name} 中 {self.vt_symbol} 瞬时利润为空，不操作")
            return

        diff = float(moment_profit['diff'])                 # 现金差
        discount = float(moment_profit['discount'])         # 折价利润
        premium = float(moment_profit['premium'])           # 溢价利润
        now = datetime.datetime.now()                       # 当前时间
        ask1 = float(moment_profit['ask1'])                 # 计算瞬时利润时的ask1
        bid1 = float(moment_profit['bid1'])                 # 计算瞬时利润时的bid1
        moment_profit_dt = datetime.datetime.strptime(moment_profit['dp_up_at'], "%Y-%m-%d %H:%M:%S")
        time_delay = (now - moment_profit_dt).seconds
        # self.print(f"延迟: {time_delay}s, 利润：{diff}, {discount}, {premium}")
        # 延迟如果太大，则不可使用
        if time_delay > 3:
            self.print(f"瞬时利润 {time_delay}s 不更新，不操作")
            return

        # 4.平仓逻辑, 未做与计算瞬时利润时价格的检验
        if self.mode == SSMode.DISCOUNT and self.etf_pos >= 1:
            self.print("赎回")
            self.redemption(volume=self.pre_ss_vol)
            return
        elif self.mode == SSMode.DISCOUNT and self.check_basket_req_bs(target_volume=0):
            self.print("卖出篮子")
            self.set_basket_target(target_volume=0)
            return
        elif self.mode == SSMode.PREMIUM and self.basket_pos >= 0.95:
            self.print("申购")
            self.purchase(volume=self.pre_ss_vol)
            return
        elif self.mode == SSMode.PREMIUM and self.pos.get(self.vt_symbol, 0) >= 100:
            self.print("卖出ETF")
            self.buy_sell_with_target(limit_price=tick.bid_price_1,
                                      target_volume=0,
                                      per_order_max=self.per_order_vol)
            return
        elif self.etf_pos == 0 and not self.check_basket_req_bs(target_volume=0):
            # 平仓结束, 可以进一步开仓， 不return
            self.mode = SSMode.NONE

        # 5.开仓逻辑
        time_now = tick.datetime.time()
        # 限制开仓时间
        if self.trade_morning_start <= time_now <= self.trade_morning_end or \
                self.trade_afternoon_start <= time_now <= self.trade_afternoon_end:
            if discount + diff >= self.discount_min_profit:         # 折价+现金差：买ETF、赎回、卖出篮子
                # 保证当前ETF tick.ask_price_1 <= ask1:
                if tick.ask_price_1 > ask1:
                    self.print(f"折价利润满足：但ETF价格已经偏离 当前 {tick.ask_price_1} 瞬时计算时 {ask1}, 不下单")
                    return
                self.print(f"折价利润: {discount + diff} >= {self.discount_min_profit} 和价格 "
                           f"{tick.ask_price_1} <= {ask1} 同时满足：买入ETF")
                if self.buy_sell_with_target(limit_price=tick.ask_price_1,
                                             target_volume=1,
                                             per_order_max=self.per_order_vol):
                    # 开仓成功
                    self.mode = SSMode.DISCOUNT
                    return

            elif premium - diff >= self.premium_min_profit:         # 溢价-现金差：买入篮子、申购、卖出ETF
                # TODO ??? 溢价是否要保证开仓 tick.bid_price_1 >= bid1? 考虑到：溢价要等买篮子、申购以后才操作ETF
                if tick.bid_price_1 < bid1:
                    self.print(f"溢价利润满足：但ETF价格已经偏离 当前 {tick.bid_price_1} 瞬时计算时 {bid1}, 不下单")
                    return
                self.print(f"溢价利润 {premium - diff} >= {self.premium_min_profit} 和价格 "
                           f"{tick.bid_price_1} >= {bid1} 同时满足：买入篮子")
                if self.set_basket_target(target_volume=1):
                    self.mode = SSMode.PREMIUM
                return
        else:
            self.print("当前时间不允许开仓")
            return

        # 6. 开仓不完全： 处理开仓只成交了一部分的情况
        if self.mode == SSMode.DISCOUNT and self.etf_pos < 1:
            self.print("ETF开仓数量不够，直接补齐")
            self.buy_sell_with_target(limit_price=tick.ask_price_1,
                                      target_volume=1,
                                      per_order_max=self.per_order_vol)
        elif self.mode == SSMode.PREMIUM and self.check_basket_req_bs(target_volume=1):
            self.print("篮子开仓数量不够，直接补齐")
            self.set_basket_target(target_volume=1)

    def on_bar(self, bar: BarData):
        """
        这里只做历史数据队列的存储
        :param bar:
        :return:
        """
        self.am.update_bar(bar)
        if not self.am.inited:
            return
        if not self.trading:
            return
        self.handle_bar(bar)

    def handle_bar(self, bar: BarData):
        """
        bar级别的策略逻辑可以写在这里
        :param bar:
        :return:
        """
        pass
        # self.cancel_all()
        # # spread = self.get_spread("IM159845-1")
        # etf_symbol = self.vt_symbol
        # dDiscount, dPremium = self.get_moment_profit(etf_symbol)
        # # print(f'cta策略中取到了 价差：{spread.name} 价格为：{spread.last_price}')
        # print(f'cta策略中取到了 {etf_symbol} 的瞬时利润： discount: {dDiscount}, premium: {dPremium}')
        #
        # if self.basket_pos < 1 and self.etf_pos == 0:
        #     print('买入ETF')
        #     self.buy_sell_with_target(limit_price=self.tick.ask_price_1,
        #                               target_volume=self.pre_ss_vol,
        #                               per_order_max=self.pre_ss_vol,
        #                               )
        #
        # elif self.etf_pos >= self.pre_ss_vol:
        #     print('赎回')
        #     self.redemption(1)
        # elif self.basket_pos >= 1:
        #     print('卖出篮子')
        #     self.set_basket_target(target_volume=0)

    def print(self, msg):
        print(f"{datetime.datetime.now()} {self.strategy_name}: {msg}")
